mod client;
mod cmd;
mod log;

use client::SetServer;
use log::{false_log, true_log, warn_log};
use cmd::{options, set_options, find_options, CmdOption };
use std::net::{TcpListener, TcpStream};
use std::io::{self, Read, Write};
use std::env::args;
use std::process::exit;
/*
fn listener() {
    // Listen mode
    let listener = TcpListener::bind(format!("{}:{}", args[2], args[3]))?;
    println!("Listening on {}:{}", args[2], args[3]);
    
    for stream in listener.incoming() {
        let mut stream = stream?;
        let mut buffer = [0; 1024];
        
        loop {
            let n = stream.read(&mut buffer)?;
            if n == 0 { break; }
            io::stdout().write_all(&buffer[..n])?;
        }
    }
    Ok(())
}
*/

fn dial_tcp(host: String, port: u16) {
    // Connect mode
    let mut stream = match TcpStream::connect((host, port)) {
        Ok(stream) => stream,
        Err(e) => {
            eprintln!("无法连接到服务器: {}", e);
            return;
        }
    };

    // read_line > stdout
    let mut stream_clone = stream.try_clone().unwrap();
    std::thread::spawn(move || {
        let mut buffer = [0; 1024];
        loop {
            match stream_clone.read(&mut buffer) {
                Ok(0) => break, 
                Ok(n) => {
                    io::stdout().write_all(&buffer[..n]).unwrap();
                    io::stdout().flush().unwrap();
                }
                Err(e) => {
                    eprintln!("读取错误: {}", e);
                    break;
                }
            }
        }
    });

    // stdin > Server
    loop {
        let mut input = String::new();
        io::stdin().read_line(&mut input).unwrap();
        if input.trim().is_empty() {
            continue;
        }
        match stream.write_all(input.as_bytes()) {
            Ok(_) => {}
            Err(e) => {
                eprintln!("发送错误: {}", e);
                break;
            }
        }
    }
}

fn main() {
    let long_options = vec![
        CmdOption {
            abbr: "-h".to_string(),
            full: "--help".to_string(),
            vari: "".to_string(),
            defa: "".to_string(),
            desc: "display this help and exit.".to_string(),
        },
        CmdOption {
            abbr: "-lp".to_string(),
            full: "--listen_port".to_string(),
            vari: "[int]".to_string(),
            defa: "false".to_string(),
            desc: "listen the local port number.".to_string(),
        },
        CmdOption {
            abbr:"-klp".to_string(),
            full:"--keep_listen_port".to_string(),
            vari:"[int]".to_string(),
            defa:"false".to_string(),
            desc:"keep to listen the local port number.".to_string(),
        },
        CmdOption {
            abbr:"-e".to_string(),
            full:"--exec".to_string(),
            vari:"[shell]".to_string(),
            defa:"".to_string(),
            desc:"program to exec after connect.".to_string(),
        },
        CmdOption {
            abbr:"-s".to_string(),
            full:"--security".to_string(),
            vari:"{0,1}".to_string(),
            defa:"0".to_string(),
            desc:"set the security mode.".to_string(),
        },
        // 添加更多选项...
    ];
    let mut args: Vec<String> = args().collect();
    if args.len() == 1 {
        help(&long_options);
        print!("\x1B[36mCmd line: \x1B[0m");
        io::stdout().flush().expect(&format!("{}Failed to flush stdout", false_log));
        let mut arg = args[0].clone() + " ";
        io::stdin().read_line(&mut arg).expect(&format!("{}Failed to read line.", false_log));
        args = arg.trim().split(" ").map(|s| s.to_string()).collect();
    }

    if set_options(&args, &long_options[0]) {
        help(&long_options);
        exit(0);
    }
    
    let mut connect = true;
    let mut s = SetServer{
        exec    :       find_options(&args, &long_options, &"-e".to_string()),
        security:       find_options(&args, &long_options, &"-s".to_string()),
        port    :       "".to_string(),
        keep    :       false
    };

    /*
    for v in 1..3 {
        if options(&args, &long_options[v]) != "false" {
            connect = false;

            if long_options[v].abbr == "-lp" {
                s.port = options(&args, &long_options[v]);
            } else if long_options[v].abbr == "-klp" {
                s.port = options(&args, &long_options[v]);
                s.keep = true;
            }
        }
    }
    */
    if connect {
        let mut connect_addr = String::new();
        let mut host = String::new();
        let mut port: u16;
        
        if args[1].contains(':') {
            let parts: Vec<&str> = args[1].split(':').collect();
            host = parts[0].to_string();
            port = parts[1].parse::<u16>().unwrap_or_else(|e| {
                eprintln!("{}{}", false_log, e);
                help(&long_options);
                exit(1);
            });
        } else if args.len() < 3 {
            println!("{}Please refer to the help for use.", false_log);
            help(&long_options);
            exit(0);
        } else {
            host = args[1].clone();
            port = args[2].parse::<u16>().unwrap_or_else(|e| {
                eprintln!("{}无法解析端口号: {}", false_log, e);
                exit(1);
            });
        }
    
        dial_tcp(host, port);
    }
}

fn help(long_options: &[CmdOption]) {
    let mut data = String::new();
    data += &format!("VY-netcat v{}, the network tools suitable for CTF.", env!("CARGO_PKG_VERSION"));
    data += "\nBasic usages:";
    data += "\n connect to somewhere:\tnc [addr] [port]";
    data += "\n \t\t\tnc [addr]:[port]";
    data += "\n listen to somewhere:\tnc -lp [port]";
    data += "\n keep to listen:\tnc -klp [port]";
    data += "\nCmdOptions:";
    println!("{}", data);

    for v in long_options {
        let mut data = format!(" {}, {} {}", v.abbr, v.full, v.vari);
        let data_len = data.chars().count();
        for _ in 0..(5 - (data_len / 8)) {
            data.push('\t');
        }
        data += &format!("{}", v.desc);
        println!("{}", data);
    }
}




